/*
   @file transpose_custom.cpp
*/

#include "kernel_operator.h"

class KernelTransDataTo5HD
{
public:
    __aicore__ inline KernelTransDataTo5HD() {}
    __aicore__ inline void Init(__gm__ uint8_t *src, __gm__ uint8_t *dstGm, __gm__ uint8_t *scratch)
    {
        int block_id = AscendC::GetBlockIdx();
        // if (block_id == 31)
        {
            AscendC::printf("the block_id is : %d\n", block_id);

            int src_offset = block_id * 16 * 512; // 行偏移
            int scr_offset = block_id * 16;       // 列偏移
            AscendC::printf("the src_offset is : %d \n the scr_offset is : %d \n", src_offset, scr_offset);

            srcGlobal.SetGlobalBuffer((__gm__ half *)src + src_offset, srcDataSize);

            __dstGlobal.SetGlobalBuffer((__gm__ half *)dstGm, srcDataSize);

            // scratchGlobal.SetGlobalBuffer((__gm__ half *)scratch + scr_offset, 16 * (512 - 16 * block_id));
            scratchGlobal.SetGlobalBuffer((__gm__ half *)scratch + scr_offset, srcDataSize);

            pipe.InitBuffer(inQueueSrc, 1, srcDataSize * sizeof(half));
            pipe.InitBuffer(outQueueDst, 1, dstDataSize * sizeof(half));
        }
    }
    __aicore__ inline void Process()
    {
        for (int i{0}; i < 16; i++)
        {
            AscendC::printf("in trans process, the i is : %d\n", i);
            CopyIn(i);
            Compute();
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int i)
    {
        AscendC::LocalTensor<half> srcLocal = inQueueSrc.AllocTensor<half>();
        AscendC::DataCopy(srcLocal, srcGlobal[i * 16], srcDataSize);

        AscendC::printf("copyin testing printf123\n");

        inQueueSrc.EnQue(srcLocal);
        inQueueSrc.FreeTensor(srcLocal);
    }
    __aicore__ inline void Compute()
    {
        AscendC::printf("testing printf123\n");

        AscendC::LocalTensor<half> srcLocal = inQueueSrc.DeQue<half>();
        AscendC::LocalTensor<half> dstLocal = outQueueDst.AllocTensor<half>();
        AscendC::TransDataTo5HDParams transDataParams;

        // AscendC::DataCopy(scratchGlobal, srcLocal, srcDataSize);

        half flag = 1.0;
        auto pos = srcLocal.GetPosition();
        AscendC::printf("the pos is :%d\n", static_cast<int>(pos));
        AscendC::printf("in compute the flag is :%f\n", flag);

        // half halfValue = srcLocal.GetValue(0);
        // int16_t value = *reinterpret_cast<int16_t *>(&halfValue);
        // AscendC::printf("after testing printf 0\n");
        // AscendC::printf("the srcLocal[0] is : %d\n", value);

        // AscendC::printf("the half of halfValue is : %f", halfValue);
        uint64_t size = srcLocal.GetSize();
        AscendC::printf("the size of srcLocal is : %d\n", size);

        transDataParams.dstHighHalf = false;
        transDataParams.srcHighHalf = false;
        transDataParams.repeatTimes = 1;
        transDataParams.dstRepStride = 0;
        transDataParams.srcRepStride = 0;
        // 入参类型是LocalTensor的调用方式
        AscendC::LocalTensor<half> dstLocalList[16];
        for (int i = 0; i < 16; i++)
        {
            dstLocalList[i] = dstLocal[width * i];
        }
        AscendC::LocalTensor<half> srcLocalList[16];
        for (int i = 0; i < 16; i++)
        {
            srcLocalList[i] = srcLocal[width * i];
        }

        // AscendC::DataCopy(scratchGlobal, srcLocal, srcDataSize);

        AscendC::TransDataTo5HD(dstLocalList, srcLocalList, transDataParams);
        // AscendC::DataCopy(scratchGlobal, dstLocalList[0], srcDataSize);

        // AscendC::DataCopy(scratchGlobal, dstLocal, srcDataSize);

        outQueueDst.EnQue<half>(dstLocal);
        inQueueSrc.FreeTensor(srcLocal);
        AscendC::printf("after testing printf 1\n");
    }
    __aicore__ inline void CopyOut(int i)
    {
        AscendC::printf("after testing printf 2\n");

        AscendC::LocalTensor<half> dstLocal = outQueueDst.DeQue<half>();

        const AscendC::SliceInfo srcSliceInfo[] = {{0, 15, 0, 1, 512}, {0, 15, 0, 1, 16}};
        const AscendC::SliceInfo dstSliceInfo[] = {{0, 15, 0, 1, 512}, {0, 15, 0, 1, 16}};
        AscendC::DataCopy(scratchGlobal[i * 512 * 16], dstLocal, dstSliceInfo, srcSliceInfo, 2);
        // AscendC::DataCopy(scratchGlobal, dstLocal, srcDataSize);
        AscendC::DataCopy(__dstGlobal, dstLocal, srcDataSize);

        outQueueDst.FreeTensor(dstLocal);
        AscendC::printf("after testing printf 3\n");
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueueSrc;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outQueueDst;

    AscendC::GlobalTensor<half> srcGlobal;
    AscendC::GlobalTensor<half> __dstGlobal;
    AscendC::GlobalTensor<half> scratchGlobal;

    int srcDataSize = 8192;
    int dstDataSize = 8192;
    int width = 512;
};

extern "C" __global__ __aicore__ void transpose_custom(GM_ADDR src, GM_ADDR dstGm, GM_ADDR scratch)
{
    if (src == nullptr || dstGm == nullptr || scratch == nullptr)
    {
        AscendC::printf("Invalid pointer in Init\n");
        return;
    }
    KernelTransDataTo5HD op;
    op.Init(src, dstGm, scratch);
    op.Process();
    AscendC::printf("kernel ending \n");
}
